from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Response
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn
import asyncio
import cv2
import base64
import numpy as np
from typing import List, Dict, Optional
import time

app = FastAPI(
    title="Robot Control API",
    description="API for controlling robots with video streaming",
    version="1.0.0"
)

# CORS配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"]
)


# 数据模型
class MoveCommand(BaseModel):
    direction: str
    session_id: str


class VoiceCommand(BaseModel):
    text: str
    session_id: str


# 机器人控制器
class RobotController:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.position = {"x": 0, "y": 0, "z": 0}
        self.speed = 1

    def move(self, direction: str):
        if direction == "forward":
            self.position["y"] += self.speed
        elif direction == "backward":
            self.position["y"] -= self.speed
        elif direction == "left":
            self.position["x"] -= self.speed
        elif direction == "right":
            self.position["x"] += self.speed
        elif direction == "up":
            self.position["z"] += self.speed
        elif direction == "down":
            self.position["z"] -= self.speed
        return self.position


# 连接管理器
class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.robots: Dict[str, RobotController] = {}
        self.lock = asyncio.Lock()

    async def connect(self, websocket: WebSocket):
        async with self.lock:
            await websocket.accept()
            session_id = str(id(websocket))
            self.active_connections[session_id] = websocket
            self.robots[session_id] = RobotController(session_id)
            return session_id

    async def disconnect(self, websocket: WebSocket):
        async with self.lock:
            session_id = str(id(websocket))
            if session_id in self.robots:
                del self.robots[session_id]
            if session_id in self.active_connections:
                del self.active_connections[session_id]


manager = ConnectionManager()


# 视频捕获器
class VideoCapture:
    def __init__(self):
        self.cap = cv2.VideoCapture(0)
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        self.cap.set(cv2.CAP_PROP_FPS, 30)
        if not self.cap.isOpened():
            self.cap = None
            self.width, self.height = 640, 480
            self.count = 0

    async def get_frame(self):
        try:
            if self.cap and self.cap.isOpened():
                ret, frame = self.cap.read()
                if not ret:
                    return None
            else:
                frame = np.zeros((self.height, self.width, 3), np.uint8)
                cv2.putText(frame, f"Simulated Video {self.count}", (10, 30),
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                self.count += 1

            _, buffer = cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70])
            return base64.b64encode(buffer).decode('utf-8')
        except Exception as e:
            print(f"Error capturing frame: {e}")
            return None


video_capture = VideoCapture()


# API路由
@app.get("/api/status")
async def get_status(session_id: str = None):
    if session_id and session_id in manager.robots:
        return {"position": manager.robots[session_id].position}
    return {"default_position": {"x": 0, "y": 0, "z": 0}}


# @app.post("/api/move")
# async def move_robot(command: MoveCommand):
#     if command.session_id in manager.robots:
#         print(1)
#         position = manager.robots[command.session_id].move(command.direction)
#         print(2)
#         return {"message": f"Moved {command.direction}", "position": position}
#     raise HTTPException(status_code=404, detail="Session not found")


@app.post("/api/move")
async def move_robot(command: MoveCommand):
    # 临时创建控制器（测试用）
    robot = RobotController(command.session_id)
    position = robot.move(command.direction)
    print(f"Move {command.direction} to {position}")  # 确认请求到达
    return {"message": f"Moved {command.direction}", "position": position}


@app.post("/api/voice")
async def process_voice(command: VoiceCommand):
    text = command.text.lower()
    direction = None

    if "forward" in text or "前进" in text:
        direction = "forward"
    elif "backward" in text or "后退" in text:
        direction = "backward"
    elif "left" in text or "左" in text:
        direction = "left"
    elif "right" in text or "右" in text:
        direction = "right"

    if direction and command.session_id in manager.robots:
        position = manager.robots[command.session_id].move(direction)
        return {"message": f"Executed voice command: {direction}", "position": position}
    return {"message": "Voice command not recognized"}


# WebSocket路由
@app.websocket("/ws/video")
async def video_stream(websocket: WebSocket):
    session_id = await manager.connect(websocket)
    last_ping = time.time()
    try:
        while True:
            # 心跳检测
            if time.time() - last_ping > 30:
                await websocket.send_text("ping")
                try:
                    pong = await asyncio.wait_for(websocket.receive_text(), timeout=5)
                    if pong != "pong":
                        raise WebSocketDisconnect()
                    last_ping = time.time()
                except:
                    raise WebSocketDisconnect()

            frame = await video_capture.get_frame()
            if frame:
                await websocket.send_text(frame)
            await asyncio.sleep(0.033)  # ~30fps
    except WebSocketDisconnect:
        await manager.disconnect(websocket)
    except Exception as e:
        print(f"WebSocket error: {e}")
        await manager.disconnect(websocket)


# 启动配置
if __name__ == "__main__":
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        ws_max_size=1024 * 1024 * 10,  # 10MB最大消息大小
        ws_ping_interval=20,  # 20秒发送一次ping
        ws_ping_timeout=30,  # 30秒无响应断开
        timeout_keep_alive=60  # 保持连接超时
    )